package cn.seecoder.ai.algorithm.ml.regression

import cn.seecoder.ai.algorithm.ml.BaseRegression
import cn.seecoder.ai.model.bo.TrainParamsBO
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.regression.RandomForestRegressor

class RandomForestRegression extends BaseRegression{

  override def buildMachineLearningStage(trainParams: TrainParamsBO): PipelineStage = {

    val rf = new RandomForestRegressor()
      .setFeaturesCol("indexedFeatures")
      .setLabelCol("label")
      .setSeed(trainParams.seed)
      .setMaxDepth(trainParams.maxDepth)
      .setMaxBins(trainParams.maxBins)
      .setMinInstancesPerNode(trainParams.minInstancesPerNode)
      .setMinInfoGain(trainParams.minInfoGain)
      .setNumTrees(trainParams.numTrees)

    rf

  }

}
